from math import e, log, sqrt
import gym
from gym import spaces
import numpy as np
import random
import os
from envs.snake import setting
import time

import copy

class SnakeEnv(gym.Env):
    def __init__(self):

        # 长和宽
        self.block_h = setting.height
        self.block_w = setting.width

        # 动作：上下左右
        self.action_space = spaces.Discrete(4)

        # 环境状态
        # 感觉没啥用
        self.observation_space = spaces.Box(low=0, high=3, shape=(self.block_h, self.block_w), dtype=np.uint8)

        # 初始/当前 的方向
        self.now_action = 2

        # 记录吃到食物的时候，蛇头和食物的距离 (似乎无效了)
        # self.distance_eat = -1

        # 记录地图的尺寸,分为small mid big,分别取值0, 1, 2
        self.map_size = 0

        # 记录障碍物的位置
        self.obstacles = []

        # 记录本轮的步数
        self.count = 0

        # 记录上一次吃到食物的步数，初始为0
        self.last_eat = 0

        # 设置两次吃到食物的最大间隔，没吃到就算死亡
        self.max_count_eat = (setting.width + setting.height) * 2



    
    # 得到奖励函数
    # 奖励设置 后期需要修改
    # 普通移动：-1
    # 吃到食物：1
    # 撞到墙和身体：-10

    def get_reward(self, flag):

        # 长度奖励
        length_reward = len(self.body) / ((setting.width - 2) * (setting.height-2))

        # 时长奖励
        step_reward = self.count / 300

        # 距离食物的距离的奖励, 靠近食物有奖励，远离食物就惩罚
        h_x, h_y = self.body[0]
        l_x, l_y = self.body[1]
        dis1 = (abs(h_x - self.food[0]) / (setting.height-2) + abs(h_y - self.food[1]) / (setting.width-2))
        dis2 = (abs(l_x - self.food[0]) / (setting.height-2) + abs(l_y - self.food[1]) / (setting.width-2))

        if dis1 < dis2:
            dis_reward = 1
        else:
            dis_reward = -1

        # 两次吃到食物的间隔奖励，越短越好
        inter_eat_reward = e ** -((self.count - self.last_eat) / (setting.width + setting.height - 2))

        # 上面可以归于移动过程的奖励

        move_reward = 0
        if flag == 'move':  # 移动奖励缩小10倍，防止蹭食物
            return (length_reward * 0.05 + step_reward * 0.05 + dis_reward * 0.2 + inter_eat_reward * 0.7) / 10
        elif flag == 'dead':
            # 不想死的话，惩罚力度大一些
            return -10
        elif flag == 'eat':
            self.last_eat = self.count
            # 吃到了东西，给大的奖励
            return 5 + length_reward * 0.5 + step_reward * 0.5

    # 生成食物
    def gen_food(self):
        f_x, f_y = random.randint(1, self.block_h-2), random.randint(1, self.block_w-2)

        def check():
            # 判断是否在蛇的身体上，在的话就重新生成
            for i in range(len(self.body)):
                if self.body[i][0] == f_x and self.body[i][1] == f_y:
                    return False
            # 判断是否在障碍物上，在的话就重新生成
            for obs in self.obstacles:
                if f_x == obs[0] and f_y == obs[1]:
                    return False
            return True
        
        while not check():
            f_x, f_y = random.randint(1, self.block_h-2), random.randint(1, self.block_w-2)
        return [f_x, f_y]

    # 生成障碍物
    def gen_obs(self):
        o_x, o_y = random.randint(1, self.block_h - 2), random.randint(1, self.block_w - 2)

        obss = []
        # 判断是否在蛇的身体上、食物和已经创建的墙上
        def check():
            # 身体的检查
            for i in range(len(self.body)):
                if self.body[i][0] == o_x and self.body[i][1] == o_y:
                    return False
            # 食物的检查
            if self.food[0] == o_x and self.food[1] == o_y:
                return False
            # 已经创建的障碍物检查
            for obs in obss:
                if obs[0] == o_x and obs[1] == o_y:
                    return False
            return True

        for _ in range(np.random.randint(setting.obs_num)):
            while not check():
                o_x, o_y = random.randint(1, self.block_h - 2), random.randint(1, self.block_w - 2)
            obss.append((o_x, o_y))
        return obss
    
    # 生成状态地图，痕迹需要遮盖
    def insert_body2state(self, cover=None):
        # 头
        self.state[self.body[0][0]][self.body[0][1]] = setting.head
        # 身体
        for i in range(1, len(self.body)):
            self.state[self.body[i][0]][self.body[i][1]] = setting.body
        # 食物
        self.state[self.food[0]][self.food[1]] = setting.food

        # 掩盖尾巴拖过的痕迹
        if cover != None:
            self.state[cover[0]][cover[1]] = setting.space
    
    def reset(self):
        # 上一次吃到食物
        self.last_eat = 0
        # 步数归0
        self.count = 0
        # 最大间隔归0
        self.max_count_eat = (setting.width + setting.height) * 2

        # 长和宽
        self.block_h = setting.height
        self.block_w = setting.width

        # 1为墙 0为空地 2为蛇头 3为蛇身 4 为食物
        self.state = np.array([[setting.space for _ in range(self.block_w)] for __ in range(self.block_h)])

        # 蛇的身体
        self.body = [(int(self.block_h / 2)-1, int(self.block_w / 2)-1), (int(self.block_h / 2)-1, int(self.block_w / 2))]

        # 食物
        self.food = self.gen_food()

        # 初始方向 ←
        self.now_action = 2

        # wall
        for i in range(self.block_h):
            self.state[i][0] = setting.wall
            self.state[i][self.block_w-1] = setting.wall
        
        for i in range(self.block_w):
            self.state[0][i] = setting.wall
            self.state[self.block_h-1][i] = setting.wall

        # 障碍物
        self.obstacles = self.gen_obs()
        for obs in self.obstacles:
            self.state[obs[0]][obs[1]] = setting.wall

        # 将上面这些插入到状态地图中
        self.insert_body2state()
        return self.get_state_list()

    def render(self):
        time.sleep(1)
        os.system('cls')
        print()
        print(self.state)


    def get_state_list(self):
        """
        将二维列表的状态转化为6个状态的列表 [dis1, dis2, f1, f2, f3, f4]
        dis1: 头的横坐标和食物的距离，dis1 / height
        dis2: 头的纵坐标和食物的距离，dis2 / width
        fi: 四个方向上是否有墙或者身体，取值为 1, 0
        """
        h_x, h_y = self.body[0]
        dis1 = (h_x - self.food[0]) / (setting.height - 2)
        dis2 = (h_y - self.food[1]) / (setting.width - 2)
        # 距离食物的位置的信息
        s = [(dis1 + dis2) / 2]
        # 两次吃到食物的间隔的信息，越大越好
        s.append((self.max_count_eat - (self.count - self.last_eat)) / ((setting.width + setting.height) * 2))

        # 判断四个方向上是否有墙或者身体
        shifts = [[-1, 0], [1, 0], [0, -1], [0, 1]]
        for shift in shifts:
            new_x, new_y = h_x + shift[0], h_y + shift[1]
            """
            越界或者没有障碍物取0，否则取1
            """
            if new_x < 0 or new_x >= setting.height or new_y < 0 or new_y >= setting.width \
                    or self.state[new_x][new_y] == setting.food or self.state[new_x][new_y] == setting.space:
                s.append(0)
            else:
                s.append(1)

        '''
        只用6元组训练出来的蛇，无法感知到全局的状态，所以容易卡死自己，见useful_model/s_1000(容易撞死自己)
        但是吧，输入全局的信息又会导致鲁棒性不高，无法适用于不同尺寸的地图，是否采用边界填充的方式?，即假设外面有很多墙
        '''
        # 采用边界填充的方式，将地图尺寸固定再20 * 20, 所以状态的大小为 20 * 20 + 6，感觉会训练很长时间

        # 计算下填充数量

        h, w = int((setting.max_size - setting.height) / 2), int((setting.max_size - setting.width) / 2)
        t = np.pad(np.array(self.state), ((h, setting.max_size - h - setting.height),
            (w, setting.max_size - w - setting.width)), 'constant', constant_values=(setting.wall)) / setting.state_num

        return s + t.reshape(-1).tolist()
        # print(len(s))
        # return s


    # 返回状态、奖励、和是否结束
    """
    在本类里面，状态一直是一个二维的列表
    返回值: next_state, reward, done, length
    """
    def step(self, action):
        # 步数 +1
        self.count += 1

        # 是否吃到食物
        iseat = False

        # 是否死亡
        isdead = False

        # 记录下蛇的尾巴，吃到食物时使用
        tail_x, tail_y = self.body[-1]
        # 挪动身体的位置（不包括头，头靠动作移动）
        for i in reversed(range(1, len(self.body))):
            self.body[i] = self.body[i-1]
        
        # 对无效的方向进行矫正 并且进行惩罚(后面又删掉了)
        self.invalid_action_penalty = 0
        if action > 3 or action < 0 or(self.now_action == 0 and action == 1) or (self.now_action == 1 and action == 0) or \
           (self.now_action == 2 and action == 3) or (self.now_action == 3 and action == 2):
            action = self.now_action
            # print('invalid action')
            self.invalid_action_penalty = -0.1
        else:
            self.now_action = action

        # 记录下头的位置
        if action == 0: # 向上
            head = (self.body[0][0] - 1, self.body[0][1])

        elif action == 1: # 向下
            head = (self.body[0][0] + 1, self.body[0][1])
 
        elif action == 2: # 向左
            head = (self.body[0][0], self.body[0][1] - 1)

        elif action == 3: # 向右
            head = (self.body[0][0], self.body[0][1] + 1)

        # 判断是否死亡 撞墙了或者撞到自己
        # print(len(self.state), len(self.state[0]), head)
        if self.state[head[0]][head[1]] == setting.body or self.state[head[0]][head[1]] == setting.wall:
            # print('死亡')
            isdead = True
        # 一定时间每吃到食物，也算死亡
        if self.count - self.last_eat > self.max_count_eat:
            isdead = True

        # 加入蛇头
        self.body[0] = head

        # 状态改变:q
        self.insert_body2state((tail_x, tail_y))

        if isdead:
            # print('game over!')
            return self.get_state_list(), self.get_reward('dead'), True, len(self.body)
        # 如果吃到东西
        if self.body[0][0] == self.food[0] and self.body[0][1] == self.food[1]:
            iseat = True
            # 蛇变长
            self.body.append((tail_x, tail_y))

            # 重新生成食物
            self.food = self.gen_food()

            # 更新两次吃到食物的间隔信息，吃到一个就可以多走4步
            self.max_count_eat += 5

            # 更新状态
            self.insert_body2state()
        if iseat:
            return self.get_state_list(), self.get_reward('eat'), False, len(self.body)
        return self.get_state_list(), self.get_reward('move'), False, len(self.body)


    """
    导出自身的数据:
    food
    state
    body
    """
    def data_export(self):
        return copy.deepcopy([self.state, self.food, self.body])


    """
    接受导出的数据
    """
    def data_import(self, d):
        self.state = d[0]
        self.food = d[1]
        self.body = d[2]

    # 改变环境的大小，在6 - 20 之间变化
    # 增加鲁棒性
    def change_map(self):
        """
        依据轮次改变地图的大小, 采用3阶段的随机生成;
        同时修改最大步数
        """
        # thresh = 0.7
        # if np.random.rand() < thresh:
        #     setting.width = np.random.randint(6,15)
        #     setting.height = np.random.randint(6,15)
        # else:
        #     setting.width = np.random.randint(15,21)
        #     setting.height = np.random.randint(15,21)

        # print(setting.height, setting.width, 'changemap')
        max_step = 0
        if self.map_size < 17:
            setting.width = np.random.randint(6, 11)
            setting.height = np.random.randint(6,11)
            max_step = 150

        elif self.map_size < 22 :
            setting.width = np.random.randint(10, 16)
            setting.height = np.random.randint(10,16)
            max_step = 180
        elif self.map_size < 25:
            setting.width = np.random.randint(16, 21)
            setting.height = np.random.randint(16, 21)
            max_step = 220

        self.map_size = (self.map_size + 1) % 25
        return max_step
     

# if __name__ == '__main__':
#     env = gym.make('Boxing-v0')
#     print(env.reset())
    # env = SnakeEnv()
    # state, done = env.reset()
    # env.render()
    # action = 0
    # while not done:
    #     m = input()
    #     if m == 'w' and action != 1:
    #         action = 0
    #     elif m == 's' and action != 0:
    #         action = 1
    #     elif m == 'a' and action != 3:
    #         action = 2
    #     elif m == 'd' and action != 2:
    #         action = 3
    #     else:
    #         action = -1
    #     if action == -1:
    #         continue            
    #     state, reward, done = env.step(action)
    #     env.render()